#include "kernel_operator.h"

constexpr int32_t BUFFER_NUM = 2;

template<typename TYPE_X, typename TYPE_Z> 
class KernelCol2Im {
    using T = TYPE_X;
public:
    __aicore__ inline KernelCol2Im() {}

    __aicore__ inline void Init(GM_ADDR col, GM_ADDR x, const Col2ImCustomTilingData &tiling_data)
    {
        this->tiling = tiling_data;

        ASSERT(AscendC::GetBlockNum() != 0 && "block dim can not be zero!");
        this->core_id = AscendC::GetBlockIdx();
        
        if (tiling.core_pixel_count[this->core_id] <= 0) {
            return;
        }

        colGm.SetGlobalBuffer((__gm__ TYPE_X *)col, tiling.input_elements);
        xGm.SetGlobalBuffer((__gm__ TYPE_Z *)x, tiling.output_elements);
        
        uint32_t elements_per_block = 32 / sizeof(TYPE_Z);
        this->aligned_tile_size = ((tiling.tile_pixel_num + elements_per_block - 1) / elements_per_block) * elements_per_block;
        
        // 初始化缓冲区
        pipe.InitBuffer(outQueueX, BUFFER_NUM, this->aligned_tile_size * sizeof(TYPE_Z));
        pipe.InitBuffer(accumBuffer, this->aligned_tile_size * sizeof(TYPE_Z));
        pipe.InitBuffer(tempBuffer, this->aligned_tile_size * sizeof(TYPE_Z));
    }
    
    __aicore__ inline void Process()
    {
        int32_t loopCount = tiling.core_loop_times[this->core_id];
        
        for (int32_t i = 0; i < loopCount; i++) 
        {
            int32_t current_tile_size;
            if (i == loopCount - 1) {
                current_tile_size = tiling.core_tail_pixels[this->core_id];
            } else {
                current_tile_size = tiling.tile_pixel_num;
            }
            
            Compute(i, current_tile_size);
            CopyOut(i, current_tile_size);
        }
    }

private:
    __aicore__ inline void Compute(int32_t progress, int32_t current_tile_size)
    {
        AscendC::LocalTensor<TYPE_Z> xLocal = outQueueX.AllocTensor<TYPE_Z>();
        AscendC::LocalTensor<TYPE_Z> accumLocal = accumBuffer.Get<TYPE_Z>();
        AscendC::LocalTensor<TYPE_Z> tempLocal = tempBuffer.Get<TYPE_Z>();
        
        // 初始化累加器为0
        TYPE_Z zero_val = static_cast<TYPE_Z>(0);
        AscendC::Duplicate(accumLocal, zero_val, this->aligned_tile_size);
        
        int32_t pixelStart = tiling.core_pixel_start[this->core_id] + progress * tiling.tile_pixel_num;
        
        // 遍历所有kernel位置进行批量累加
        for (int32_t kh = 0; kh < tiling.kernel_h; kh++) {
            for (int32_t kw = 0; kw < tiling.kernel_w; kw++) {
                // 初始化临时buffer
                AscendC::Duplicate(tempLocal, zero_val, this->aligned_tile_size);
                
                // 批量收集当前kernel位置对所有像素的贡献
                for (int32_t i = 0; i < current_tile_size; i++) {
                    int32_t pixelIdx = pixelStart + i;
                    
                    if (pixelIdx >= tiling.total_output_pixels) {
                        continue;
                    }
                    
                    // 解析像素坐标
                    int32_t n, c, h, w;
                    PixelIndexToNCHW(pixelIdx, n, c, h, w);
                    
                    // 检查该kernel位置是否对当前像素有贡献
                    int32_t h_out, w_out;
                    if (!CalculateOutputPosition(h, w, kh, kw, h_out, w_out)) {
                        continue;
                    }
                    
                    // 计算col索引
                    int32_t colIdx = CalculateColIndex(n, c, kh, kw, h_out, w_out);
                    
                    if (colIdx >= 0 && colIdx < tiling.input_elements) {
                        TYPE_X value = colGm.GetValue(colIdx);
                        tempLocal.SetValue(i, static_cast<TYPE_Z>(value));
                    }
                }
                
                // 批量累加：accumLocal += tempLocal
                AscendC::Add(accumLocal, accumLocal, tempLocal, this->aligned_tile_size);
            }
        }
        
        // 复制结果
        AscendC::DataCopy(xLocal, accumLocal, this->aligned_tile_size);
        
        tempBuffer.FreeTensor(tempLocal);
        accumBuffer.FreeTensor(accumLocal);
        outQueueX.EnQue<TYPE_Z>(xLocal);
    }

    __aicore__ inline void CopyOut(int32_t progress, int32_t current_tile_size)
    {
        AscendC::LocalTensor<TYPE_Z> xLocal = outQueueX.DeQue<TYPE_Z>();
        
        int32_t pixelOffset = tiling.core_pixel_start[this->core_id] + progress * tiling.tile_pixel_num;
        
        if (current_tile_size > 0) {
            uint32_t bytes_to_copy = current_tile_size * sizeof(TYPE_Z);
            
            if (bytes_to_copy % 32 == 0) {
                AscendC::DataCopy(xGm[pixelOffset], xLocal, current_tile_size);
            } else {
                AscendC::DataCopyExtParams copyParams{
                    1,
                    bytes_to_copy,
                    0,
                    0,
                    0
                };
                
                AscendC::DataCopyPad(xGm[pixelOffset], xLocal, copyParams);
            }
        }
        
        outQueueX.FreeTensor(xLocal);
    }

    __aicore__ inline void PixelIndexToNCHW(int32_t pixelIdx, int32_t& n, int32_t& c, int32_t& h, int32_t& w)
    {
        int32_t HW = tiling.H * tiling.W;
        int32_t CHW = tiling.C * HW;
        
        n = pixelIdx / CHW;
        int32_t remaining = pixelIdx % CHW;
        c = remaining / HW;
        remaining = remaining % HW;
        h = remaining / tiling.W;
        w = remaining % tiling.W;
    }

    __aicore__ inline bool CalculateOutputPosition(int32_t h, int32_t w, int32_t kh, int32_t kw, 
                                                     int32_t& h_out, int32_t& w_out)
    {
        int32_t kh_dilated = kh * tiling.dilation_val;
        int32_t kw_dilated = kw * tiling.dilation_val;
        
        int32_t h_numerator = h + tiling.padding_val - kh_dilated;
        int32_t w_numerator = w + tiling.padding_val - kw_dilated;
        
        if (h_numerator % tiling.stride_val != 0 || w_numerator % tiling.stride_val != 0) {
            return false;
        }
        
        h_out = h_numerator / tiling.stride_val;
        w_out = w_numerator / tiling.stride_val;
        
        if (h_out < 0 || h_out >= tiling.out_H || w_out < 0 || w_out >= tiling.out_W) {
            return false;
        }
        
        return true;
    }

    __aicore__ inline int32_t CalculateColIndex(int32_t n, int32_t c, int32_t kh, int32_t kw, 
                                                  int32_t h_out, int32_t w_out)
    {
        int32_t channel_block = c * (tiling.kernel_h * tiling.kernel_w) + kh * tiling.kernel_w + kw;
        int32_t pos_idx = h_out * tiling.out_W + w_out;
        int32_t colIdx = n * (tiling.input_channels * tiling.L) + channel_block * tiling.L + pos_idx;
        
        return colIdx;
    }

private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueX;
    AscendC::TBuf<AscendC::QuePosition::VECCALC> accumBuffer;
    AscendC::TBuf<AscendC::QuePosition::VECCALC> tempBuffer;
    AscendC::GlobalTensor<TYPE_X> colGm;
    AscendC::GlobalTensor<TYPE_Z> xGm;

    int32_t core_id;
    int32_t aligned_tile_size;

    Col2ImCustomTilingData tiling;
};

extern "C" __global__ __aicore__ void col2_im_custom(GM_ADDR col, GM_ADDR x, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling); 
    KernelCol2Im<DTYPE_COL, DTYPE_X> op;
    op.Init(col, x, tiling_data);
    op.Process();
}